import os
import tempfile
import fitz  # PyMuPDF
from PIL import Image
import pytesseract
from pypdf import PdfReader
import pdfplumber
import pandas as pd
import re
from typing import Dict, List, Optional, Tuple
import requests

class PDFProcessor:
    def __init__(self):
        self.temp_files = []
        
    def __del__(self):
        """析构函数，清理临时文件"""
        self.cleanup_temp_files()
    
    def cleanup_temp_files(self):
        """清理所有临时文件"""
        for temp_file in self.temp_files:
            try:
                if os.path.exists(temp_file):
                    os.remove(temp_file)
            except Exception as e:
                print(f"删除临时文件失败 {temp_file}: {e}")
        self.temp_files = []
    
    def is_pdf_corrupted(self, pdf_path: str) -> bool:
        """
        检测PDF文件是否损坏
        返回: True表示损坏，False表示正常
        """
        try:
            # 方法1: 使用PyPDF2检测基本结构
            with open(pdf_path, 'rb') as file:
                reader = PdfReader(file)
                if len(reader.pages) == 0:
                    return True
                
                # 尝试读取第一页内容
                first_page = reader.pages[0]
                text = first_page.extract_text()
                
            # 方法2: 使用PyMuPDF进一步验证
            doc = fitz.open(pdf_path)
            doc.close()
            
            return False
        except Exception as e:
            print(f"PDF文件损坏检测异常: {e}")
            return True
    
    def repair_pdf(self, pdf_path: str) -> Optional[str]:
        """
        尝试修复损坏的PDF文件
        返回: 修复后的文件路径，如果修复失败返回None
        """
        try:
            # 创建临时文件用于修复
            temp_dir = tempfile.gettempdir()
            repaired_path = os.path.join(temp_dir, f"repaired_{os.path.basename(pdf_path)}")
            self.temp_files.append(repaired_path)
            
            # 方法1: 使用pikepdf尝试修复（这里简化实现）
            # 实际应用中可以使用pikepdf库进行更复杂的修复
            with open(pdf_path, 'rb') as f_in:
                pdf_data = f_in.read()
            
            # 检查PDF文件头尾
            if b'%PDF-' not in pdf_data[:1024]:
                # 添加缺失的PDF头
                pdf_data = b'%PDF-1.4\n' + pdf_data
            
            with open(repaired_path, 'wb') as f_out:
                f_out.write(pdf_data)
            
            # 验证修复结果
            if not self.is_pdf_corrupted(repaired_path):
                return repaired_path
            else:
                return None
                
        except Exception as e:
            print(f"PDF修复失败: {e}")
            return None
    
    def is_scanned_pdf(self, pdf_path: str, threshold: float = 0.1) -> bool:
        """
        判断PDF是否为扫描版
        基于文本区域占比和可提取文本量判断
        """
        try:
            doc = fitz.open(pdf_path)
            total_text_chars = 0
            sample_pages = min(5, len(doc))  # 抽样前5页
            
            for i in range(sample_pages):
                page = doc[i]
                
                # 提取文本
                text = page.get_text()
                total_text_chars += len(text.strip())
                
                # 检查文本块数量
                text_blocks = page.get_text("blocks")
                text_block_count = len([b for b in text_blocks if len(b[4].strip()) > 0])
                
                # 如果文本块很少但页面有内容，可能是扫描版
                if text_block_count < 2 and page.rect.area > 0:
                    doc.close()
                    return True
            
            doc.close()
            
            # 如果平均每页文本字符数很少，可能是扫描版
            avg_chars_per_page = total_text_chars / sample_pages
            return avg_chars_per_page < 50  # 阈值可根据需要调整
            
        except Exception as e:
            print(f"扫描版检测失败: {e}")
            return True  # 出错时按扫描版处理
    
    def extract_text_with_pypdf(self, pdf_path: str) -> str:
        """使用pypdf提取文本版PDF的文本内容"""
        text_content = ""
        try:
            with open(pdf_path, 'rb') as file:
                reader = PdfReader(file)
                for page in reader.pages:
                    text = page.extract_text()
                    if text:
                        text_content += text + "\n\n"
        except Exception as e:
            print(f"pypdf文本提取失败: {e}")
        return text_content
    
    def extract_tables_with_pypdf(self, pdf_path: str) -> str:
        """使用pdfplumber提取表格（markdown格式）"""
        tables_md = ""
        try:
            with pdfplumber.open(pdf_path) as pdf:
                for page_num, page in enumerate(pdf.pages):
                    tables = page.extract_tables()
                    if tables:
                        for i, table in enumerate(tables):
                            if table and any(any(cell is not None for cell in row) for row in table):
                                tables_md += f"### 第{page_num + 1}页表格{i + 1}\n\n"
                                
                                # 转换为DataFrame便于处理
                                df = pd.DataFrame(table)
                                
                                # 清理空行空列
                                df = df.dropna(how='all').dropna(axis=1, how='all')
                                
                                if not df.empty:
                                    # 转换为markdown表格
                                    tables_md += df.to_markdown(index=False) + "\n\n"
        except Exception as e:
            print(f"表格提取失败: {e}")
        return tables_md
    
    def extract_links_with_pypdf(self, pdf_path: str) -> str:
        """提取超链接"""
        links_info = "## 超链接\n\n"
        try:
            doc = fitz.open(pdf_path)
            for page_num in range(len(doc)):
                page = doc[page_num]
                links = page.get_links()
                
                if links:
                    for link in links:
                        if 'uri' in link:
                            links_info += f"- [{link['uri']}]({link['uri']}) (第{page_num + 1}页)\n"
            doc.close()
        except Exception as e:
            print(f"超链接提取失败: {e}")
        
        return links_info if links_info != "## 超链接\n\n" else ""
    
    def extract_images_with_pypdf(self, pdf_path: str) -> Tuple[Dict[str, Image.Image], str]:
        """提取图片并返回PIL图像对象和markdown描述"""
        images_dict = {}
        images_md = "## 图片\n\n"
        image_count = 0
        
        try:
            doc = fitz.open(pdf_path)
            for page_num in range(len(doc)):
                page = doc[page_num]
                image_list = page.get_images()
                
                for img_index, img in enumerate(image_list):
                    try:
                        xref = img[0]
                        pix = fitz.Pixmap(doc, xref)
                        
                        if pix.n - pix.alpha < 4:  # GRAY or RGB
                            img_data = pix.tobytes("ppm")
                        else:  # CMYK
                            pix_cmyk = fitz.Pixmap(fitz.csRGB, pix)
                            img_data = pix_cmyk.tobytes("ppm")
                            pix_cmyk = None
                        
                        # 转换为PIL Image
                        pil_img = Image.open(io.BytesIO(img_data))
                        
                        image_id = f"page{page_num + 1}_img{img_index + 1}"
                        images_dict[image_id] = pil_img
                        
                        images_md += f"### {image_id}\n"
                        images_md += f"- 位置: 第{page_num + 1}页\n"
                        images_md += f"- 尺寸: {pil_img.size}\n"
                        images_md += f"- 模式: {pil_img.mode}\n\n"
                        
                        image_count += 1
                        pix = None
                        
                    except Exception as e:
                        print(f"图片{img_index}处理失败: {e}")
                        continue
                        
            doc.close()
        except Exception as e:
            print(f"图片提取失败: {e}")
        
        return images_dict, images_md if image_count > 0 else ""
    
    def ocr_pdf_text(self, pdf_path: str) -> str:
        """使用OCR提取扫描版PDF的文本"""
        ocr_text = ""
        try:
            doc = fitz.open(pdf_path)
            for page_num in range(len(doc)):
                page = doc[page_num]
                
                # 将页面转换为图像
                pix = page.get_pixmap()
                img_data = pix.tobytes("ppm")
                pil_img = Image.open(io.BytesIO(img_data))
                
                # 使用Tesseract进行OCR
                text = pytesseract.image_to_string(pil_img, lang='chi_sim+eng')
                if text.strip():
                    ocr_text += f"## 第{page_num + 1}页\n\n{text}\n\n"
                
            doc.close()
        except Exception as e:
            print(f"OCR文本提取失败: {e}")
        
        return ocr_text
    
    def ocr_pdf_tables(self, pdf_path: str) -> str:
        """使用OCR尝试识别表格（简化版）"""
        # 注意：这是一个简化的表格识别，实际应用中可能需要更复杂的表格检测算法
        tables_md = "## 表格识别结果\n\n"
        tables_md += "> 注意：扫描版PDF的表格识别可能不准确，建议人工核对\n\n"
        
        try:
            doc = fitz.open(pdf_path)
            for page_num in range(min(3, len(doc))):  # 只处理前3页以提高效率
                page = doc[page_num]
                pix = page.get_pixmap()
                img_data = pix.tobytes("ppm")
                pil_img = Image.open(io.BytesIO(img_data))
                
                # 使用Tesseract检测可能包含表格的区域
                # 这里使用简单的文本布局分析，实际应用可能需要更复杂的算法
                ocr_data = pytesseract.image_to_data(pil_img, output_type=pytesseract.Output.DICT)
                
                # 简单的表格检测逻辑（根据文本对齐方式）
                # 实际应用中可以考虑使用专门的表格识别库
                tables_md += f"### 第{page_num + 1}页可能的表格区域\n\n"
                tables_md += "表格识别功能需要进一步优化...\n\n"
                
            doc.close()
        except Exception as e:
            print(f"OCR表格识别失败: {e}")
        
        return tables_md
    
    def process_pdf(self, pdf_path: str) -> Dict:
        """
        处理PDF文件的主函数
        返回包含解析结果的字典
        """
        result = {
            'metadata': {},
            'text': '',
            'tables': '',
            'links': '',
            'images': {},
            'images_md': '',
            'full_markdown': '',
            'is_scanned': False,
            'was_repaired': False
        }
        
        try:
            # 1. 检查PDF是否损坏
            if self.is_pdf_corrupted(pdf_path):
                print("PDF文件损坏，尝试修复...")
                repaired_path = self.repair_pdf(pdf_path)
                if repaired_path:
                    pdf_path = repaired_path
                    result['was_repaired'] = True
                    print("PDF修复成功")
                else:
                    print("PDF修复失败，尝试继续处理...")
            
            # 2. 判断是否为扫描版
            is_scanned = self.is_scanned_pdf(pdf_path)
            result['is_scanned'] = is_scanned
            print(f"PDF类型: {'扫描版' if is_scanned else '文本版'}")
            
            # 3. 根据类型采用不同方法解析
            if is_scanned:
                # 扫描版使用OCR处理
                result['text'] = self.ocr_pdf_text(pdf_path)
                result['tables'] = self.ocr_pdf_tables(pdf_path)
                
                # 扫描版也可以尝试提取图片和链接
                images_dict, images_md = self.extract_images_with_pypdf(pdf_path)
                result['images'] = images_dict
                result['images_md'] = images_md
                result['links'] = self.extract_links_with_pypdf(pdf_path)
            else:
                # 文本版使用常规方法
                result['text'] = self.extract_text_with_pypdf(pdf_path)
                result['tables'] = self.extract_tables_with_pypdf(pdf_path)
                result['links'] = self.extract_links_with_pypdf(pdf_path)
                
                images_dict, images_md = self.extract_images_with_pypdf(pdf_path)
                result['images'] = images_dict
                result['images_md'] = images_md
            
            # 4. 组合成完整的markdown文档
            result['full_markdown'] = self.generate_markdown(result)
            
            # 5. 获取元数据
            try:
                with open(pdf_path, 'rb') as file:
                    reader = PdfReader(file)
                    if hasattr(reader, 'metadata') and reader.metadata:
                        result['metadata'] = dict(reader.metadata)
            except:
                pass
            
        except Exception as e:
            print(f"PDF处理过程中发生错误: {e}")
            result['error'] = str(e)
        
        finally:
            # 清理临时文件
            self.cleanup_temp_files()
        
        return result
    
    def generate_markdown(self, result: Dict) -> str:
        """生成完整的markdown文档"""
        markdown_parts = []
        
        # 添加标题和元数据
        markdown_parts.append("# PDF解析结果\n\n")
        markdown_parts.append(f"**文档类型**: {'扫描版' if result['is_scanned'] else '文本版'}\n")
        if result.get('was_repaired'):
            markdown_parts.append("**状态**: 已修复\n")
        markdown_parts.append("\n")
        
        # 添加文本内容
        if result['text']:
            markdown_parts.append("## 文本内容\n\n")
            markdown_parts.append(result['text'])
            markdown_parts.append("\n")
        
        # 添加表格
        if result['tables']:
            markdown_parts.append(result['tables'])
            markdown_parts.append("\n")
        
        # 添加图片描述
        if result['images_md']:
            markdown_parts.append(result['images_md'])
            markdown_parts.append("\n")
        
        # 添加超链接
        if result['links']:
            markdown_parts.append(result['links'])
            markdown_parts.append("\n")
        
        return "".join(markdown_parts)


# 使用示例
def main():
    import io  # 用于内存字节流操作
    
    processor = PDFProcessor()
    
    # 示例PDF文件路径（请替换为实际路径）
    pdf_path = "example.pdf"
    
    if not os.path.exists(pdf_path):
        print(f"文件 {pdf_path} 不存在")
        return
    
    try:
        # 处理PDF
        result = processor.process_pdf(pdf_path)
        
        # 输出结果
        print("=" * 50)
        print("PDF处理完成!")
        print(f"文档类型: {'扫描版' if result['is_scanned'] else '文本版'}")
        print(f"是否修复: {'是' if result.get('was_repaired', False) else '否'}")
        print(f"提取图片数量: {len(result['images'])}")
        print(f"Markdown总字符数: {len(result['full_markdown'])}")
        print("=" * 50)
        
        # 保存markdown结果到文件
        output_md = "pdf_result.md"
        with open(output_md, 'w', encoding='utf-8') as f:
            f.write(result['full_markdown'])
        print(f"Markdown结果已保存到: {output_md}")
        
        # 显示前500个字符作为预览
        print("\n内容预览:")
        print(result['full_markdown'][:500] + "..." if len(result['full_markdown']) > 500 else result['full_markdown'])
        
    except Exception as e:
        print(f"处理失败: {e}")
    
    finally:
        # 确保清理临时文件
        processor.cleanup_temp_files()


if __name__ == "__main__":
    main()